from qrsfilter import QRSFilter
import numpy as np
import struct
import matplotlib.pyplot as plt
from scipy.signal import butter, lfilter, freqz
import wfdb
from wfdb import processing
import os


class qrsdetector:
    # QRS检测器
    def __init__(self):
        self.qrsfilter = QRSFilter()
        # 阈值
        self.Thr = 20
        self.st = 0
        self.peak = 0  # 峰值
        self.SincePeak = 0
        self.filterDelay = 44

    def do(self, ecg):
        x = self.qrsfilter.filter(ecg)
        if self.st == 0:
            if x > self.Thr:
                self.st = 1
                self.peak = x
                self.SincePeak = 0

        elif self.st == 1:
            if x > self.peak:
                self.SincePeak = 0
                self.peak = x
            self.SincePeak = self.SincePeak + 1
            if x < self.Thr:
                self.st = 0
                return x, self.SincePeak, True
        return x, 0, False


def qrs_file(path, record):
    '''
    提取文件的QRS波位置，目前针对MIT格式文件，
    输入：
        path 文件路径
        record 文件名称，不包括后缀
    输出：
        rpos:R波的位置
        raw：原始的心电数据
        filtx：通过QRSFilter后的滤波数据
    '''
    fp = open(os.path.join(path, record + '.dat'), "rb")
    detector = qrsdetector()
    rpos = []
    filtx = []
    rawx = []
    nsample = 0
    while True:
        d = fp.read(2)
        if not d:
            break
        ecg = struct.unpack('h', d)
        x, delay, isQrs = detector.do(ecg[0])
        filtx.append(x)
        rawx.append(ecg)
        if isQrs:
            rpos.append(nsample - delay-44)
        nsample = nsample + 1

    fp.close()

    return rpos, rawx, filtx

# 评估
def qrs_eval(path, record, rpos):
    annotation = wfdb.rdann(os.path.join(path, record), 'atr')#读取.art文件
    anntyp = np.array(annotation.symbol)
    ref = annotation.sample[anntyp != '+']

    comparitor = processing.compare_annotations(
        ref_sample=ref, test_sample=np.array(rpos), window_width=int(0.1 * 250))
    res = [comparitor.tp, comparitor.fn, comparitor.fp]
    return res, ref


def dirfile(path, ext):
    """遍历文件夹下某一类型后的所有文件
       # path:路径
       # ext:后缀
    """
    filelist = []
    for record in os.listdir(path):  # 遍历整个文件夹
        if os.path.splitext(record)[1] == ".{}".format(ext):
            filelist.append(os.path.splitext(record)[0])
    filelist = np.array(filelist)
    filelist.sort()
    return filelist

if __name__ == '__main__':
    path = "D:/wuhui/testcode/SmartHealth-master/DATA/mitdb/"
    recordList = dirfile(path, 'atr')
    resAll = []
    m = 1
    for record in recordList:
        rpos, rawx, filtx = qrs_file(path, record)
        res, ref = qrs_eval(path, record, rpos)
        print(record, res)
        resAll.append(res)
        m = m + 1
        if m > 48:
            break
    m = np.array(resAll).sum(axis=0)
    sen = m[0] / (m[0] + m[1]) #敏感度
    ppv = m[0] / (m[0] + m[2]) #精确率/阳性预测值
    print(m)
    print('sen = {:.4f}; ppv = {:.4f}'.format(sen, ppv))

    plt.subplot(2, 1, 1)
    plt.plot(filtx)
    plt.plot(rpos, np.array(filtx)[rpos], '.')
    plt.subplot(2, 1, 2)
    plt.plot(rawx)
    plt.show()






